// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/testing/expression/arg_types_generator.h>

namespace kumo::pollux::fuzzer {

/// An argument type generator for decimal function. A map keyed on the pair of
/// precision and scale could be initialized with all possible input types.
/// Argument types are generated by looking up the map with the precision and
/// scale of return type, and randomly selecting valid input types. Derived
/// classes should call 'initialize' from the constructor and specify the number
/// of decimal arguments. They should also implement toReturnType with matching
/// number of pairs of precision and scale.
class DecimalArgTypesGeneratorBase : public ArgTypesGenerator {
 public:
  std::vector<TypePtr> generateArgs(
      const exec::FunctionSignature& signature,
      const TypePtr& returnType,
      FuzzerGenerator& rng) override;

 protected:
  // Computes result type for all possible pairs of decimal input types. Stores
  // the results in 'inputs_' map keyed by the precision and scale of return
  // type.
  // @param numArgs the number of decimal argument types. It only supports
  // initialization with one or two argument types.
  virtual void initialize(uint32_t numArgs);

  // Given precisions and scales of the inputs, returns precision and scale of
  // the result. Returns std::nullopt if a valid return type cannot be generated
  // with inputs. Used when the return type is generated with one pair of input
  // precision and scale.
  virtual std::optional<std::pair<int, int>> toReturnType(int p, int s) {
    POLLUX_UNREACHABLE();
  }

  // Used when the return type is generated with two pairs of input precision
  // and scale.
  virtual std::optional<std::pair<int, int>>
  toReturnType(int p1, int s1, int p2, int s2) {
    POLLUX_UNREACHABLE();
  }

 private:
  // Returns randomly selected pair of input types that produce the specified
  // result type.
  std::vector<TypePtr> findInputs(
      const TypePtr& returnType,
      FuzzerGenerator& rng) const;

  // Maps from the precision and scale of return type to corresponding input
  // types.
  melon::F14FastMap<std::pair<int, int>, std::vector<std::vector<TypePtr>>>
      inputs_;
};

} // namespace kumo::pollux::fuzzer
